 library(bayesSurv); #library(panel);
#library(MCMCpack);  
 library(mvtnorm);  library(lattice)
 library(msm); library(MASS);  library(Matrix)
# library(accuracy);
#library(kinship)

  ###############################
  Index<-function(Names){
   uniq <- unique(na.omit(Names))                # get units
   J <- length(uniq)                          # how many units 
   index <- rep (NA, J)                      # how many observations for each unit
   for (i in 1:J) index[Names==uniq[i]] <- i
   index
 }
 remove.NA <- function (x) 
{
    x[complete.cases(x), ]
}

  Wishart <- function (element, col){
    mm <- array(0, c(col,col))
    mm[lower.tri(mm, diag = TRUE)] <-element
    res <- mm + t(mm)
    diag(res) <- diag(res)/2
    return (res)
  }

  sumtrans<-function(Mat){
    total<-0
    for (r in 1: nrow(Mat)){
      total<-total+outerself(Mat[r,])
    }
    total
  }
 outerself<-function(x) outer(x,x)
 Covariance.update <- function(Covariate, prior.freedom, prior.scale,
                               dimension, obs.length, current.value){
   if (ncol(Covariate)==1){                                    
     Ei<-rgamma(1,(prior.freedom+obs.length)/2,
                rate=(prior.scale+sum(current.value^2))/2)
     Ei <- as.matrix(Ei)
     return(list(Ei=Ei))
   }else{
     OEin <- solve(prior.scale)+sumtrans(Mat=current.value)
     sE<-rWishart (1, prior.freedom+obs.length, solve(OEin))
     Ei <- Wishart(element=sE, col=dimension)
  }
     return(list(Ei=Ei, sE=sE))
 }

 GroupPred <- function(X, index){
   b<-ncol(X)
   J <- length(unique(na.omit(index)))
   XPred<-matrix(0, nrow=J, ncol=b)
   for (i in 1:b){
     Y<-X[,i]
     a <- length(Y)
     Z <- is.na(X[,i])==TRUE
     m <- which(Z==TRUE)
     s <- length(Z[which(Z==FALSE)])
     if (s==a){
       XP <- tapply(Y, index, mean)
       XPred[,i]<-XP
     }else{
       XP <- tapply(X[-m, i], index[-m], mean)
       XPred[,i]<-XP
    }
   }
   return(XPred)
 }

  ###################################
  random <- function(Var, Pred, index){
    J<-length(unique(na.omit(index)))
    RanFix <- matrix(0, nrow=1, ncol=ncol(Var)*ncol(Pred))
    for (j in 1:J){
      W <- as.matrix(Var[which(index==j),])
      a <- kronecker(Diagonal(ncol(Var)), t(as.matrix(Pred[j,])))
      RanFix<-rbind(RanFix, as.matrix(W%*%a))
    }
    TTT<-RanFix[-1,]
  }
################################
 randomt2 <- function(Var, Pred, index1, index2, entry){
   index <- na.omit(index1)
   index2 <- na.omit(index2)
   index2 <- index2-min(index2)+1
   Var <- Var
   Pred <- na.omit(Pred)
   J1<-length(unique(na.omit(index)))
   WWW <- matrix(NA, nrow=length(entry), ncol=ncol(Var)*ncol(Pred))
   for (j in index){
     rows <- entry[which(index==j)]
     TT <- index2[rows]
     W <- as.matrix(Var[rows,])
     PW <- as.matrix(Pred[TT,])
     Ee <- length(rows)
     Tt <- matrix(NA, ncol=ncol(Var)*ncol(Pred), nrow=Ee)
     for (s in 1:Ee){
       Tt[s,] <- as.vector(t(cbind(W[s,])%*%rbind(PW[s,])))
     }
     WWW[rows,] <- Tt    
   }
   WWW
 }

 RegressionData <- function(yob, Unit, Time, Fixed, UnitRandom, TimeRandom,
                             UnitPred, TimePred){
    if(!is.numeric(UnitPred)& !is.numeric(TimePred)){ # no predictor, only random intercept or coefficients
      NewData <- DataArrange.NP(id1=Unit, id2=Time, y=yob, Fixed=Fixed,
                                UnitRandom=UnitRandom, TimeRandom=TimeRandom)
    } else{
        if (!is.numeric(UnitPred) &is.numeric(TimePred)){ # only time-specifc predictors
           NewData <- DataArrange.TP(id1=Unit, id2=Time, y=yob, Fixed=Fixed,
                                     UnitRandom=UnitRandom, TimeRandom=TimeRandom,
                                     TimePred=TimePred)
        } else{
           if (is.numeric(UnitPred) &!is.numeric(TimePred)){ # only unit-specific predictors
              NewData <- DataArrange.UP(id1=Unit, id2=Time, y=yob, Fixed=Fixed,
                                        UnitRandom=UnitRandom, TimeRandom=TimeRandom,
                                        UnitPred=UnitPred)
             }else{ # both predictors
               NewData <- DataArrange(id1=Unit, id2=Time, y=yob, Fixed=Fixed,
                                      UnitRandom=UnitRandom, TimeRandom=TimeRandom,
                                      UnitPred=UnitPred, TimePred=TimePred)
             }
        }
   }
    return(NewData)
  }

 ###################################
 DataArrange <- function(id1, id2, y, Fixed, UnitRandom,
                         TimeRandom, UnitPred, TimePred){
  # Make sure the design matrices are matrix formats
   if(!is.matrix(Fixed)){Fixed <- as.matrix(Fixed)}
   if(!is.matrix(UnitRandom)){UnitRandom <- as.matrix(UnitRandom)}
   if(!is.matrix(TimeRandom)){TimeRandom <- as.matrix(TimeRandom)}
   if(!is.matrix(UnitPred)){UnitPred <- as.matrix(UnitPred)}
   if(!is.matrix(TimePred)){TimePred <- as.matrix(TimePred)}

   J<-length(unique(na.omit(id1)))  # how many units, in case NAs
   len <- 1:length(na.omit(id1))    # how many observations in the dataset

   #_____  Create group-level predictors____________#

   # create a design matrix  at the group level of units
   UnitP <-GroupPred(X=as.matrix(UnitPred), index=id1)
   A <- as.matrix(UnitP)      
   # create a design matrix  at the group level of units
   UnitQ <-GroupPred(X=as.matrix(TimePred), index=id2) 
   B <- as.matrix(UnitQ)   

   #_______ Interacting terms in the individual level ____#

   # interacting terms of unit-level variables and individual
   # variables--- they are with fixed effects in the reduced model
   RanFix <- random(Var=UnitRandom, Pred=A, index=id1)

   # interacting terms of time-level variables and individual
   #  variables--- they are with fixed effects in the reduced model
   TimFix <- randomt2(Var=TimeRandom, Pred=B,
                      index1=id1, index2=id2, entry=len)

   #________ Get the dimensions ___________________________#

   RD <- ncol(RanFix)
   TD <- ncol(TimFix)
   FD <- ncol(Fixed)
   UD <- ncol(UnitRandom)
   UT <- ncol(TimeRandom)

   #________ Creat the dataset ______________________________#

   # conbined the data with two indices
   RawData<-cbind(id1, id2, y, Fixed, RanFix, TimFix, UnitRandom, TimeRandom) 
   Data2 <- remove.NA(RawData)

   # get rid of entries with missing data---no drop-out:
   # THIS NEEDS FURTHER THINKING---HOW TO DEAL WITH MISSING DATA

   DD <- ncol(Data2)
   #__________ Recale the data and get right covariates ___________#
   newid1 <- Data2[,1]
   newid2 <- Data2[,2]
   Y<-Data2[,3]
   X <- as.matrix(cbind(Data2[, c(4: (3+FD+TD+RD))]))
   W <- as.matrix(Data2[,  c((4 +FD+TD+RD):(3+FD+TD+RD+UD))])
   S <- as.matrix(Data2[, (4+FD+TD+RD+UD):(3+FD+TD+RD+UD+UT)])

   return(list(newid1, newid2,Y, X, W, S))
  }

 ###################################
 DataArrange.UP <- function(id1, id2, y, Fixed,
                             UnitRandom, TimeRandom, UnitPred){
   # Make sure the design matrices are matrix formats
   if(!is.matrix(Fixed)){Fixed <- as.matrix(Fixed)}
   if(!is.matrix(UnitRandom)){UnitRandom <- as.matrix(UnitRandom)}
   if(!is.matrix(TimeRandom)){TimeRandom <- as.matrix(TimeRandom)}
   if(!is.matrix(UnitPred)){UnitPred <- as.matrix(UnitPred)}

   J<-length(unique(na.omit(id1)))  # how many units, in case NAs
   len <- 1:length(na.omit(id1))    # how many observations in the dataset

   #______ Create group-level predictors______________#

   # create a design matrix  at the group level of units
   UnitP <-GroupPred(X=as.matrix(UnitPred), index=id1)
   A <- as.matrix(UnitP)      
   #_____ Interacting terms in the individual level _____#
   # interacting terms of unit-level variables and individual
   # variables--they are with fixed effects in the reduced model
   RanFix <- random(Var=UnitRandom, Pred=A, index=id1)
   #_______ Get the dimensions _________________________#

   RD <- ncol(RanFix)
   FD <- ncol(Fixed)
   UD <- ncol(UnitRandom)
   UT <- ncol(TimeRandom)

   #__________ Creat the dataset _______________________#
   # conbined the data with two indices
   if (ncol(TimeRandom)==1&length(unique(TimeRandom))==1){
     RawData<-cbind(id1, id2, y, Fixed,  RanFix, UnitRandom)
   }else{
     RawData<-cbind(id1, id2, y, Fixed, RanFix, UnitRandom, TimeRandom)
   }
   Data2 <- remove.NA(RawData)

   # get rid of entries with missing data---no drop-out:
   # THIS NEEDS FURTHER THINKING---HOW TO DEAL WITH MISSING DATA 
   DD <- ncol(Data2)
   #______ Recale the data and get right covariates ___#

   newid1 <- Data2[,1]
   newid2 <- Data2[,2]
   Y<-Data2[,3]
   X <- as.matrix(cbind(Data2[, c(4: (3+FD+RD))]))
   W <- as.matrix(Data2[, c((4+FD+RD):(3+FD+RD+UD))])
   if (ncol(TimeRandom)==1&length(unique(TimeRandom))==1){
     S <- as.matrix(rep(1, nrow(X)))
   }else{
     S <- as.matrix(Data2[, c((4+FD+RD+UD):(3+FD+RD+UD+UT))])
   }
   return(list(newid1, newid2,Y, X, W, S))
 }

 ###################################
 DataArrange.TP <- function(id1, id2, y, Fixed, UnitRandom,
                            TimeRandom,  TimePred){
   # Make sure the design matrices are matrix formats
   if(!is.matrix(Fixed)){Fixed <- as.matrix(Fixed)}
   if(!is.matrix(UnitRandom)){UnitRandom <- as.matrix(UnitRandom)}
   if(!is.matrix(TimeRandom)){TimeRandom <- as.matrix(TimeRandom)}
   if(!is.matrix(TimePred)){TimePred <- as.matrix(TimePred)}

   J<-length(unique(na.omit(id1)))  # how many units, in case NAs
   len <- 1:length(na.omit(id1))    # how many observations in the dataset

   #__________ Create group-level predictors____________#
   # create a design matrix  at the group level of units
   UnitQ <-GroupPred(X=as.matrix(TimePred), index=id2) 
   B <- as.matrix(UnitQ)   

   #______ Interacting terms in the individual level _____#
   # interacting terms of time-level variables and individual
   # variables--- they are with fixed effects in the reduced model
   TimFix <- randomt2(Var=TimeRandom, Pred=B,
                      index1=id1, index2=id2, entry=len)

   #__________ Get the dimensions _________________#

   TD <- ncol(TimFix)
   FD <- ncol(Fixed)
   UD <- ncol(UnitRandom)
   UT <- ncol(TimeRandom)

   #________ Creat the dataset ____________________#

   # conbined the data with two indices
     RawData<-cbind(id1, id2, y, Fixed, UnitRandom, TimeRandom, TimFix)
   Data2 <- remove.NA(RawData)

   # get rid of entries with missing data---no drop-out:
   # THIS NEEDS FURTHER THINKING---HOW TO DEAL WITH MISSING DATA
   DDT <- ncol(Data2)
   #__________ Recale the data and get right covariates _____#
   newid1 <- Data2[,1]
   newid2 <- Data2[,2]
   Y<-Data2[,3]
   X <- as.matrix(cbind(Data2[, -c(c(1:3), c((DDT-UT-UD-TD+1):(DDT-TD)))]))
   if (ncol(UnitRandom)==1&length(unique(UnitRandom))==1){
     W <- as.matrix(rep(1, nrow(X)))
   }else{
     W <-  as.matrix(Data2[, (4+FD):(3+FD+UD)])
   }
   S <- as.matrix(Data2[,  c((DDT-UT+1):DDT)])
   return(list(newid1, newid2,Y, X, W, S))
 }

 # Function: there is no group level predictors, if only varing-intercepts
 # are included, the universal intercept should not be included.
 DataArrange.NP <- function(id1, id2, y, Fixed, UnitRandom, TimeRandom){
   if(!is.matrix(Fixed)){Fixed <- as.matrix(Fixed)}
   if(!is.matrix(UnitRandom)){UnitRandom <- as.matrix(UnitRandom)}
   if(!is.matrix(TimeRandom)){TimeRandom <- as.matrix(TimeRandom)}

   # Make sure the design matrices are matrix formats
   J<-length(unique(na.omit(id1)))  # how many units, in case NAs
   len <- 1:length(na.omit(id1))    # how many obsedrvations in the dataset

   FD <- ncol(Fixed)
   UD <- ncol(UnitRandom)
   UT <- ncol(TimeRandom)
   
   #_________ Creat the dataset ____________#
   # conbined the data with two indices
   if ((UD==1&length(unique(UnitRandom))==1)&(UT==1&length(unique(TimeRandom))==1)){
     RawData<-cbind(id1, id2, y, Fixed)
     Data <- remove.NA(RawData)
     newid1 <- Data[,1]
     newid2 <- Data[,2]
     Y<-Data[,3]
     X <- as.matrix(Data[, -c(1:3)])
     W <- as.matrix(rep(1,nrow(X)))
     S <- as.matrix(rep(1,nrow(X)))
   }else{
     if ((UD==1&length(unique(UnitRandom))==1)&(UT!=1|length(unique(TimeRandom))!=1)){
       RawData<-cbind(id1, id2, y, Fixed,  TimeRandom)
       Data <- remove.NA(RawData)
       RR <- ncol(Data)
       newid1 <- Data[,1]
       newid2 <- Data[,2]
       Y<-Data[,3]
       X <- as.matrix(Data[, -c(1:3)])
       W <- as.matrix(rep(1,nrow(X)))
       S <- as.matrix(Data[, (RR-UT+1):RR])
      }else{
        if ((UD!=1|length(unique(UnitRandom))!=1)&(UT==1&length(unique(TimeRandom))==1)){
          RawData<-cbind(id1, id2, y, Fixed, UnitRandom)
          Data <- remove.NA(RawData)
          RR <- ncol(Data)
          newid1 <- Data[,1]
          newid2 <- Data[,2]
          Y <- Data[,3]
          X <- as.matrix(Data[, -c(1:3)])
          W <- as.matrix(Data[, (RR-UD+1):RR])
          S <- as.matrix(rep(1,nrow(X)))
         }else{
          RawData<-cbind(id1, id2, y, Fixed, UnitRandom, TimeRandom)
          Data <- remove.NA(RawData)
          RR <- ncol(Data)
          newid1 <- Data[,1]
          newid2 <- Data[,2]
          Y<-Data[,3]
          X <- as.matrix(Data[, -c(1:3)])
          W <- as.matrix(Data[, (RR-UD -UT+1):(RR-UT)])
          S <- as.matrix(Data[, (RR-UT+1):RR])
        }
      }
   }
   return(list(newid1, newid2,Y, X, W, S))
 }


RFGenerator <- function(y,X1, Wi, St, Ai, Ft, Unit.Index, Time.Index,
                            timeint.add, unitint.add){
     NewData <- RegressionData(yob=y, Unit=Unit.Index, Time=Time.Index, Fixed=X1,
                               UnitRandom=Wi, TimeRandom=St,
                               UnitPred=Ai, TimePred=Ft)

    #Get results from the approprate function
    newid1 <- NewData[[1]]
    newid2 <- NewData[[2]]
    Y <- NewData[[3]]
    X <- as.matrix(NewData[[4]])
    W <- as.matrix(NewData[[5]])
    if (unitint.add==TRUE){
     W <- cbind(1, W)
    }
     S <- as.matrix(NewData[[6]])
    if (timeint.add==TRUE){
     S <- cbind(1, S)
    } 
    return(list(unit.id=newid1, time.id=newid2, response=Y, fixed.X=X, unit.random.W=W,
                time.random.S=S))
   }

     # several functions
     logdeter <- function(eigen.t, rho.t){
       core <- 1-rho.t*eigen.t
       log.deter <- sum(log(core[which(core>0)]))
       return(log.deter)
     }
